{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-20T15:38:56.496417600Z",
     "start_time": "2023-12-20T15:38:55.925411900Z"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "import torch\n",
    "import numpy\n",
    "from dataset_extractor import *\n",
    "from sklearn.metrics import f1_score\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('data/bert-base-chinese')\n",
    "model_cls = AutoModelForSequenceClassification.from_pretrained('data/bert-base-legal-chinese-epoch-3', num_labels=202)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-20T14:23:51.797349600Z",
     "start_time": "2023-12-20T14:23:51.081262500Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "27415 lines have been extracted\n",
      "content:公诉机关指控3月28日20时许被告人颜某在本市洪山区马湖新村足球场马路边捡拾到被害人谢某的VIVOX5手机一部并在同年3月28日21时起分多次通过支付宝小额免密支付功能秘密盗走被害人谢某支付宝内人民币3723元案发后被告人颜某家属已赔偿被害人全部损失并取得谅解公诉机关认为被告人颜某具有退赃取得谅解自愿认罪等处罚情节建议判处被告人颜某一年以下××××或者××并处罚金\n",
      "label:1\n",
      "articles:[264]\n",
      "punish_of_money:1000\n",
      "criminals:['颜某']\n",
      "death_penalty:False\n",
      "imprisonment:4\n",
      "life_imprisonment:False\n"
     ]
    }
   ],
   "source": [
    "test_pair = extract(\"data/dataset/cail/exercise_contest/data_test.json\")\n",
    "pair = {\"content\": [], \"label\": [], 'articles': [], 'punish_of_money': [], \"criminals\": [], \"death_penalty\": [],\n",
    "        \"imprisonment\": [], \"life_imprisonment\": []}\n",
    "for key in pair.keys():\n",
    "    print(f\"{key}:{test_pair[key][0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-20T14:54:07.969200Z",
     "start_time": "2023-12-20T14:53:58.151095200Z"
    }
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "batch_size = 50\n",
    "test_batch = tokenizer(test_pair[\"content\"], max_length=512, truncation=True, padding=\"max_length\", return_tensors=\"pt\")\n",
    "test = TensorDataset(test_batch[\"input_ids\"], test_batch[\"attention_mask\"], torch.tensor(test_pair[\"label\"]))\n",
    "test_dataloader = DataLoader(test, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-20T16:18:42.110191300Z",
     "start_time": "2023-12-20T16:18:42.084973Z"
    }
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cpu\")\n",
    "\n",
    "def test(model):\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "    y_predict = torch.empty(-1).to(device)\n",
    "    y_true = torch.empty(-1).to(device)\n",
    "\n",
    "    print(y_predict.shape)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for step, batch in enumerate(test_dataloader):\n",
    "            if step % 10 == 0 and not step == 0:\n",
    "                print(\"step: \", step)\n",
    "            b_input_ids = batch[0].to(device)\n",
    "            b_input_mask = batch[1].to(device)\n",
    "            b_labels = batch[2].to(device)\n",
    "\n",
    "            outputs = model(b_input_ids,\n",
    "                            token_type_ids=None,\n",
    "                            attention_mask=b_input_mask,\n",
    "                            labels=b_labels)\n",
    "\n",
    "            print(torch.argmax(outputs.logits, dim=1).shape)\n",
    "            y_predict = torch.cat((y_predict, torch.argmax(outputs.logits, dim=1).to(device)), dim=1)\n",
    "            y_true = torch.cat((y_true, b_labels), dim=1)\n",
    "\n",
    "            y_true = y_true.cpu().numpy().astype(int)\n",
    "            y_predict = y_predict.cpu().numpy().astype(int)\n",
    "            # print(y_true)\n",
    "            # print(y_predict)\n",
    "            print(f1_score(y_predict, y_true, labels=range(202), average=\"micro\"))\n",
    "            break\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-20T16:18:42.460814300Z",
     "start_time": "2023-12-20T16:18:42.418236800Z"
    }
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Trying to create tensor with negative dimension -1: [-1]",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mRuntimeError\u001B[0m                              Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[204], line 1\u001B[0m\n\u001B[1;32m----> 1\u001B[0m test(model_cls)\n",
      "Cell \u001B[1;32mIn[203], line 6\u001B[0m, in \u001B[0;36mtest\u001B[1;34m(model)\u001B[0m\n\u001B[0;32m      4\u001B[0m model\u001B[38;5;241m.\u001B[39mto(device)\n\u001B[0;32m      5\u001B[0m model\u001B[38;5;241m.\u001B[39meval()\n\u001B[1;32m----> 6\u001B[0m y_predict \u001B[38;5;241m=\u001B[39m torch\u001B[38;5;241m.\u001B[39mempty(\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m)\u001B[38;5;241m.\u001B[39mto(device)\n\u001B[0;32m      7\u001B[0m y_true \u001B[38;5;241m=\u001B[39m torch\u001B[38;5;241m.\u001B[39mempty(\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m)\u001B[38;5;241m.\u001B[39mto(device)\n\u001B[0;32m      9\u001B[0m \u001B[38;5;28mprint\u001B[39m(y_predict\u001B[38;5;241m.\u001B[39mshape)\n",
      "\u001B[1;31mRuntimeError\u001B[0m: Trying to create tensor with negative dimension -1: [-1]"
     ]
    }
   ],
   "source": [
    "test(model_cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-12-20T15:35:34.652571900Z",
     "start_time": "2023-12-20T15:35:34.646953900Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "cail",
   "language": "python",
   "display_name": "cail"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
